import pandas as pd
import matplotlib.pyplot as plt

from utils.LinearRegression.LinearRegressionUtils import LinearRegression

# 读取数据
data = pd.read_csv('../static/data/world-happiness-report-2017.csv')

# 选择训练(0.8)和测试数据(0.2)
train_data = data.sample(frac=0.8, random_state=200)
test_data = data.drop(train_data.index)

# 选择特征和目标变量
input_param_name = "Economy..GDP.per.Capita."
output_param_name = "Happiness.Score"
# 将输入参数名对应的训练数据转换为numpy数组，并进行reshape操作
x_train = train_data[input_param_name].values.reshape(-1, 1)
y_train = train_data[output_param_name].values.reshape(-1, 1)
x_test = test_data[input_param_name].values.reshape(-1, 1)
y_test = test_data[output_param_name].values.reshape(-1, 1)

# 图例展示数据分布,不同颜色显示数据集
plt.scatter(x_train, y_train, color='red')
plt.scatter(x_test, y_test, color='blue')
plt.title('Happiness Score vs Economy')
plt.xlabel('Economy')
plt.ylabel('Happiness Score')
plt.legend(['train', 'test'])
plt.show()


# 训练模型
linear_regression = LinearRegression(y_train,x_train)
# 定义训练所需参数
alpha = 0.01
num_iterations = 500
cost_history = linear_regression.train(alpha, num_iterations)

# 输出训练初始和结束时的损失值对比
print("Initial loss:", cost_history[0])
print("Final loss:", cost_history[-1])

# 绘制损失值变化图
plt.plot(range(num_iterations), cost_history)
plt.title('Cost History')
plt.xlabel('Iteration')
plt.ylabel('Cost')
plt.show()

# 绘制预测值变化直线图
plt.scatter(x_train, y_train, color='red')
plt.scatter(x_test, y_test, color='blue')
plt.plot(x_train, linear_regression.get_predict(x_train), color='green')
plt.title('Happiness Score vs Economy')
plt.xlabel('Economy')
plt.ylabel('Happiness Score')
plt.legend(['train', 'test', 'prediction'])
plt.show()










